import requests
import json

class AnythingLLM:
    def __init__(self, key):
        self.key = key

    def get_spaces(self):
        headers = {
            'Authorization':f'Bearer {self.key}',
            'accept': 'application/json'
        }

        response = requests.get('http://localhost:3001/api/v1/workspaces', headers=headers)

        if response.status_code == 200:
            return json.loads(response.content)
        else:
            return None
        
    def creat_space(self, name):
        headers = {
            'Authorization':f'Bearer {self.key}',
            'accept': 'application/json',
        }

        data = {
            'name': name,
            "similarityThreshold": 0.7,
            "openAiTemp": 0.7,
            "openAiHistory": 20,
            "openAiPrompt": "根据以下对话、相关背景和后续问题，回答用户当前提出的问题。根据需要，按照用户说明，仅返回您对上述信息问题的回答。",
            "queryRefusalResponse": "此工作区中没有相关信息来回答您的查询。",
            "chatMode": "chat",
            "topN": 4
        }

        response = requests.post('http://localhost:3001/api/v1/workspaces/new', headers=headers, json=data)

        if response.status_code == 200:
            return json.loads(response.content)
        else:
            return None
        
    def automate_knowledge_upload(self, space_id, file_path):
        id = self.upload_knowledge(space_id, file_path)
        self.link_to_workspace(id, space_id)

    def upload_knowledge(self, space_id, file_path):
        url = 'http://localhost:3001/api/v1/document/upload'

        headers = {
            'Authorization': f'Bearer {self.key}',
            'accept': 'application/json'
        }

        files = {'file': open(file_path, 'rb')}

        response = requests.post(url, headers=headers, files=files)

        if response.status_code == 200:
            return response.json()['documents'][0]['location']
        else:
            raise Exception(f"上传失败: {response.text}")
        
    def link_to_workspace(self, document_id, workspace_id):
        url = f'http://localhost:3001/api/v1/workspace/{workspace_id}/update-embeddings'
        headers = {
            'Authorization': f'Bearer {self.key}',
            'Content-Type': 'application/json'
        }
        payload = {
            "adds": [
                document_id
            ]
        }

        response = requests.post(url, headers=headers, json=payload)
        if response.status_code != 200:
            raise Exception(f"关联工作区失败: {response.text}")

